from fastapi import FastAPI, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer
from jose import JWTError, jwt
import redis
from datetime import datetime, timedelta
from typing import Optional

# 初始化FastAPI、Redis、JWT配置
app = FastAPI(title="RAG System API")
redis_client = redis.Redis(host="localhost", port=6379, db=0)
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
SECRET_KEY = "your-secret-key"  # 生产环境需用环境变量存储
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 30

# JWT鉴权依赖函数
def get_current_user(token: str = Depends(oauth2_scheme)):
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise credentials_exception
    except JWTError:
        raise credentials_exception
    return username

# Redis限流装饰器
def rate_limit(key: str, limit: int = 100, period: int = 60):
    def decorator(func):
        async def wrapper(*args, **kwargs):
            current = redis_client.incr(key)
            if current == 1:
                redis_client.expire(key, period)
            if current > limit:
                raise HTTPException(
                    status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                    detail="Too many requests, please try again later"
                )
            return await func(*args, **kwargs)
        return wrapper
    return decorator

# 示例接口：用户请求RAG服务（需鉴权+限流）
@app.post("/rag/query")
@rate_limit(key=lambda request, current_user: f"user:{current_user}")  # 按用户限流
async def rag_query(query: str, current_user: str = Depends(get_current_user)):
    # 后续调用检索层与生成层逻辑
    return {"user": current_user, "query": query, "status": "processing"}